from abc import ABC, abstractmethod
from typing import Dict, Any
import torch


class TrainingInterface(ABC):
    """Interface for training strategies.

    This defines the contract for all training implementations,
    making them interchangeable in the learner architecture.
    """

    @abstractmethod
    def train_task(
        self,
        network: torch.nn.Module,
        train_loader: torch.utils.data.DataLoader,
        test_loader: torch.utils.data.DataLoader,
        args: Dict[str, Any],
        device: torch.device,
        *largs,
        **kwargs,
    ) -> Dict[str, Any]:
        """Train the network on the current task.

        Args:
            network: The neural network to train
            train_loader: DataLoader for training data
            test_loader: DataLoader for testing data
            args: Configuration arguments
            **kwargs: Additional arguments

        Returns:
            Dictionary containing training results
        """
        pass

    @abstractmethod
    def _train_epoch(
        self,
        network: Any,
        train_loader: Any,
        criterion: Any,
        optimizer: Any,
        device: Any,
        *largs,
        **kwargs,
    ) -> Dict[str, Any]:
        """Train for one epoch.

        Args:
            network: The neural network to train
            train_loader: DataLoader for training data
            criterion: Loss function
            optimizer: Optimizer
            *largs: Additional positional arguments
            **kwargs: Additional keyword arguments

        Returns:
            Dictionary containing epoch results
        """
        pass


# FIXME implementation unclear
class ClassifierAlignmentInterface(ABC):
    """Interface for classifier alignment strategies.

    Defines the contract for aligning classifier weights or decision boundaries
    across incremental learning tasks to mitigate catastrophic forgetting.
    """

    @abstractmethod
    def register_hooks(self, learner: Any) -> None:
        """Register hooks for the alignment process.

        Args:
            learner: The learner instance to register hooks with
        """
        pass

    @abstractmethod
    def align(
        self, current_model: Any, previous_model: Any, device: Any, **kwargs
    ) -> None:
        """Align current model's classifier with previous model.

        Args:
            current_model: Current task's model
            previous_model: Previous task's model
            device: Computation device
            **kwargs: Additional parameters for alignment
        """
        pass

    @abstractmethod
    def compute_alignment_statistics(
        self, current_model: Any, previous_model: Any, data_loader: Any, device: Any
    ) -> Any:
        """Compute statistics needed for classifier alignment.

        Args:
            current_model: Current task's model
            previous_model: Previous task's model
            data_loader: DataLoader for evaluation data
            device: Computation device

        Returns:
            Statistics used for alignment analysis
        """
        pass

    @abstractmethod
    def apply_alignment(
        self,
        model: Any,
        alignment_params: Any,
        known_classes: int,
        total_classes: int,
        device: Any,
    ) -> Any:
        """Apply computed alignment to the model.

        Args:
            model: Model to apply alignment to
            alignment_params: Parameters for alignment
            known_classes: Number of classes from previous tasks
            total_classes: Total number of classes including current task
            device: Computation device

        Returns:
            Model with aligned classifier
        """
        pass


class DriftCompensationInterface(ABC):
    """Interface for drift compensation strategies.

    This defines the contract for all drift compensation implementations,
    making them interchangeable in the learner architecture.
    """

    @abstractmethod
    def compensate(
        self,
        network: Any,
        old_network: Any,
        train_loader: Any,
        class_means: Any,
        known_classes: int,
        device: Any,
        **kwargs,
    ) -> None:
        """Apply drift compensation to align feature spaces.

        Args:
            network: Current network model
            old_network: Previous task's network model
            train_loader: DataLoader for training data
            class_means: Class prototype means
            known_classes: Number of previously known classes
            device: Computation device
            **kwargs: Additional arguments
        """
        pass
